{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import symforce\n",
    "\n",
    "symforce.set_symbolic_api(\"sympy\")\n",
    "symforce.set_log_level(\"warning\")\n",
    "\n",
    "import symforce.symbolic as sf\n",
    "from symforce.ops import StorageOps\n",
    "from symforce.ops import LieGroupOps\n",
    "\n",
    "epsilon = 1e-9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tangent_D_storage(a):\n",
    "    \"\"\"\n",
    "    Computes the jacobian of the storage space of an element with respect to the tangent space around\n",
    "    that element.\n",
    "    \"\"\"\n",
    "    # Perturb a in the storage space\n",
    "    storage_dim = LieGroupOps.storage_dim(a)\n",
    "    xi = sf.M(storage_dim, 1).symbolic(\"xi\")\n",
    "    storage_perturbed = sf.M(LieGroupOps.to_storage(a)) + xi\n",
    "    a_perturbed = LieGroupOps.from_storage(a, storage_perturbed)\n",
    "    a_perturbed_tangent = sf.M(LieGroupOps.local_coordinates(a, a_perturbed))\n",
    "\n",
    "    # Compute jacobian of storage wrt perturbation\n",
    "    tangent_D_storage = a_perturbed_tangent.jacobian(xi)\n",
    "\n",
    "    # Evaluate at perturbation == zero\n",
    "    tangent_D_storage = tangent_D_storage.subs(xi, xi.zero())\n",
    "\n",
    "    return tangent_D_storage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tangent_D_storage_approx(a, epsilon):\n",
    "    \"\"\"\n",
    "    Computes the jacobian of the storage space of an element with respect to the tangent space around\n",
    "    that element.\n",
    "\n",
    "    This is an approximation - note that the exact jacobian can often be recovered with a call to\n",
    "    nsimplify with the appropriate tolerance (though this requires the use of sympy rather than symengine)\n",
    "    \"\"\"\n",
    "    # Perturb a in the storage space\n",
    "    storage_dim = LieGroupOps.storage_dim(a)\n",
    "    xi = sf.M(storage_dim, 1).symbolic(\"xi\")\n",
    "    storage_perturbed = sf.M(LieGroupOps.to_storage(a)) + xi\n",
    "    a_perturbed = LieGroupOps.from_storage(a, storage_perturbed)\n",
    "    a_perturbed_tangent = sf.M(LieGroupOps.local_coordinates(a, a_perturbed))\n",
    "\n",
    "    # Compute jacobian of storage wrt perturbation\n",
    "    tangent_D_storage = a_perturbed_tangent.jacobian(xi)\n",
    "\n",
    "    # Rather than computing the limit, we substitude a small value for xi to approximate the limit\n",
    "    # NOTE: This is much faster than taking the limit in sympy, but returns an approximation of the true\n",
    "    # jacobian.\n",
    "    assert epsilon != 0\n",
    "    tangent_D_storage = tangent_D_storage.subs(xi, epsilon * xi.one())\n",
    "\n",
    "    return tangent_D_storage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rot2 = sf.Rot2.symbolic(\"A\")\n",
    "display(tangent_D_storage_approx(rot2, epsilon))\n",
    "display(tangent_D_storage(rot2))\n",
    "display(tangent_D_storage(rot2).subs(rot2.z.squared_norm(), 1).subs(rot2.z.squared_norm(), 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rot3 = sf.Rot3().symbolic(\"A\")\n",
    "display(tangent_D_storage(rot3))\n",
    "(\n",
    "    tangent_D_storage(rot3)\n",
    "    .subs(sf.Min(1, rot3.q.squared_norm()), rot3.q.squared_norm())\n",
    "    .subs(sf.Max(0, 1 - rot3.q.squared_norm() ** 2), 1 - rot3.q.squared_norm() ** 2)\n",
    "    .subs(rot3.q.squared_norm(), sf.Symbol(\"v\"))\n",
    "    .limit(sf.Symbol(\"v\"), 1)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pose2 = sf.Pose2.symbolic(\"A\")\n",
    "pose2_tangent_D_storage = pose2.storage_D_tangent().mat.pinv()\n",
    "display(pose2_tangent_D_storage)\n",
    "pose2_tangent_D_storage.subs(pose2.R.z.squared_norm(), 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pose3 = sf.Pose3.symbolic(\"A\")\n",
    "pose3_tangent_D_storage = pose3.storage_D_tangent().mat.pinv()\n",
    "\n",
    "# This takes a while\n",
    "simplified = sf.simplify(pose3_tangent_D_storage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "more_simplified = simplified.subs(pose3.R.q.squared_norm(), 1)\n",
    "more_simplified"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The bottom right block is equal to R^{-1}, which is just R^T\n",
    "sf.simplify(pose3.R.to_rotation_matrix().matrix_inverse().mat) - more_simplified[3:, 4:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For Pose3, if we don't want to wait for sympy, we can also get the result ourselves, because we know that Pose3's storage_D_tangent has the form:\n",
    "$$\n",
    "{}_S{D}_T^P = \\begin{bmatrix}\n",
    "{}_S{D}_T^R & 0 \\\\\n",
    "0 & R\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "where ${}_S{D}_T^R$ is Rot3's storage_D_tangent, and $R$ is the rotation matrix for Rot3.  Because ${}_S{D}_T^P$ has linearly independent columns, its pseudoinverse can be computed as $A^+ = (A^T A)^{-1} A^T$ (see https://en.wikipedia.org/wiki/Moore%E2%80%93Penrose_inverse).  Writing this out,\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "{}_T{D}_S^P &= \\left( {}_S{D}_T^P \\right)^+ \\\\\n",
    "&= \\left(\n",
    "\\begin{bmatrix}\n",
    "({}_S{D}_T^R)^T & 0 \\\\\n",
    "0 & R^T\n",
    "\\end{bmatrix}\n",
    "\\begin{bmatrix}\n",
    "{}_S{D}_T^R & 0 \\\\\n",
    "0 & R\n",
    "\\end{bmatrix}\n",
    "\\right)^{-1}\n",
    "\\begin{bmatrix}\n",
    "({}_S{D}_T^R)^T & 0 \\\\\n",
    "0 & R^T\n",
    "\\end{bmatrix} \\\\\n",
    "&= \\begin{bmatrix}\n",
    "({}_S{D}_T^R)^T {}_S{D}_T^R & 0 \\\\\n",
    "0 & R^T R\n",
    "\\end{bmatrix}^{-1}\n",
    "\\begin{bmatrix}\n",
    "({}_S{D}_T^R)^T & 0 \\\\\n",
    "0 & R^T\n",
    "\\end{bmatrix} \\\\\n",
    "&= \\begin{bmatrix}\n",
    "\\frac{1}{4}({}_S{D}_T^R)^+ {}_S{D}_T^R & 0 \\\\\n",
    "0 & \\mathbb{1}\n",
    "\\end{bmatrix}^{-1}\n",
    "\\begin{bmatrix}\n",
    "\\frac{1}{4}({}_S{D}_T^R)^+ & 0 \\\\\n",
    "0 & R^T\n",
    "\\end{bmatrix} \\\\\n",
    "&= \\begin{bmatrix}\n",
    "\\frac{1}{4} \\mathbb{1} & 0 \\\\\n",
    "0 & \\mathbb{1}\n",
    "\\end{bmatrix}^{-1}\n",
    "\\begin{bmatrix}\n",
    "\\frac{1}{4} {}_T{D}_S^R & 0 \\\\\n",
    "0 & R^T\n",
    "\\end{bmatrix} \\\\\n",
    "&= \\begin{bmatrix}\n",
    "4 \\mathbb{1} & 0 \\\\\n",
    "0 & \\mathbb{1}\n",
    "\\end{bmatrix}\n",
    "\\begin{bmatrix}\n",
    "\\frac{1}{4} {}_T{D}_S^R & 0 \\\\\n",
    "0 & R^T\n",
    "\\end{bmatrix} \\\\\n",
    "&= \\begin{bmatrix}\n",
    "{}_T{D}_S^R & 0 \\\\\n",
    "0 & R^T\n",
    "\\end{bmatrix} \\\\\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
